import numpy as np
from sklearn.preprocessing import OrdinalEncoder
from deepctr.feature_column import SparseFeat, DenseFeat, get_feature_names
from sklearn.preprocessing import MinMaxScaler
from tensorflow.keras.initializers import RandomNormal


class Embedding:
    def __init__(self, feature_names, cat_feats, dense_features):
        self.feature_names = feature_names
        self.cat_feats = cat_feats
        self.dense_features = dense_features
        self.embed_features = None
        self.mms = MinMaxScaler(feature_range=(0, 1))
        self.ode = OrdinalEncoder()

    def embed_fit(self, fit_data):
        train_x = fit_data.copy()
        train_x[self.cat_feats] = train_x[self.cat_feats].fillna('-1', )
        train_x[self.dense_features] = train_x[self.dense_features].fillna(0, )

        self.cat_feats = [cat_feat for cat_feat in self.cat_feats if cat_feat in train_x.columns]
        self.dense_features = [dense_feat for dense_feat in self.dense_features if dense_feat in train_x.columns]

        # 1.Label Encoding for sparse features,and do simple Transformation for dense features
        self.ode.fit(train_x[self.cat_feats])
        if self.dense_features:
            self.mms.fit(train_x[self.dense_features])

        # 2.count #unique features for each sparse field,and record dense feature field name
        # keras.initializers.glorot_normal(seed=2022))
        sparse_feats = [SparseFeat(feat, vocabulary_size=train_x[feat].nunique(), embedding_dim=4,
                        embeddings_initializer=RandomNormal(mean=0.0, stddev=0.0001, seed=2022)) for feat in self.cat_feats]
        dense_feats = [DenseFeat(feat, 1, ) for feat in self.dense_features]
        fixlen_feature_columns = sparse_feats + dense_feats

        dnn_feature_columns = fixlen_feature_columns
        linear_feature_columns = fixlen_feature_columns
        self.embed_features = get_feature_names(linear_feature_columns + dnn_feature_columns)
        return linear_feature_columns, dnn_feature_columns

    def embed_transform(self, trans_data):
        trans_data[self.cat_feats] = trans_data[self.cat_feats].fillna('-1', )
        trans_data[self.dense_features] = trans_data[self.dense_features].fillna(0, )

        # 3.generate input data for model
        trans_category = self.ode.transform(trans_data[self.cat_feats])
        trans_data[self.cat_feats] = trans_category[0].astype(np.float64, copy=False)

        if self.dense_features:
            trans_data[self.dense_features] = self.mms.transform(trans_data[self.dense_features])
        model_input = {name: trans_data[name] for name in self.embed_features}
        return model_input